import numpy as np
import os


def right_shift_with_rounding(x, bits):
    if bits == 0:
        return x
    x1 = x >> (bits - 1)
    is_even = (x1 & 1) == 0
    x_even = x1 >> 1
    x_odd = (x1 + 1) >> 1
    result = np.where(is_even, x_even, x_odd)
    return result

def activation(input_data, bias, mult, B_BIT=16, A_BIT=8):
    N_SAMPLES, N_CH = input_data.shape
    MAXOUT = (1 << A_BIT) - 1
    output_data = np.zeros((N_SAMPLES, N_CH), dtype=np.int32)
    for n in range(N_SAMPLES):
        for ch in range(N_CH):
            o = input_data[n, ch].astype(np.int32) + bias[ch].astype(np.int32)
            mult_result = o * mult[ch].astype(np.int32)
            shifted = right_shift_with_rounding(mult_result, B_BIT)
            activated = np.clip(shifted, 0, MAXOUT)
            output_data[n, ch] = activated
    
    return output_data

def save_input_to_mem(input_data, filename, mode, P_CH, B_BIT=16):
    N_SAMPLES, N_CH = input_data.shape
    FOLD = N_CH // P_CH
    with open(filename, mode) as f:
        for n in range(N_SAMPLES):
            for fold_idx in range(FOLD):
                hex_values = []
                for i in range(P_CH):
                    ch_idx = fold_idx * P_CH + i
                    val = int(input_data[n, ch_idx])
                    # Convert to unsigned representation for hex
                    if val < 0:
                        val = val & ((1 << B_BIT) - 1)
                    hex_width = B_BIT // 4
                    hex_values.append(f"{val:0{hex_width}x}")
                # Write in little-endian order (reverse)
                f.write(''.join(reversed(hex_values)) + '\n')
    
    print(f"Saved input to {filename}")


def save_params_to_mem(params, filename, mode, N_CH, P_CH, BIT_WIDTH):
    FOLD = N_CH // P_CH
    with open(filename, mode) as f:
        for fold_idx in range(FOLD):
            hex_values = []
            for i in range(P_CH):
                ch_idx = fold_idx * P_CH + i
                val = int(params[ch_idx])
                # Convert to unsigned representation for hex
                if val < 0:
                    val = val & ((1 << BIT_WIDTH) - 1)
                hex_width = BIT_WIDTH // 4
                hex_values.append(f"{val:0{hex_width}x}")
            # Write in little-endian order (reverse)
            f.write(''.join(reversed(hex_values)) + '\n')
    
    print(f"Saved parameters to {filename}")


def save_output_to_mem(output, filename, mode, P_CH, A_BIT=8):
    """Save output data to memory file, reorganized for hardware with P_CH parallelism"""
    N_SAMPLES, N_CH = output.shape
    FOLD = N_CH // P_CH
    
    with open(filename, mode) as f:
        for n in range(N_SAMPLES):
            for fold_idx in range(FOLD):
                hex_values = []
                for i in range(P_CH):
                    ch_idx = fold_idx * P_CH + i
                    val = int(output[n, ch_idx])
                    # Convert to unsigned representation for hex
                    if val < 0:
                        val = val & ((1 << A_BIT) - 1)
                    hex_width = A_BIT // 4
                    hex_values.append(f"{val:0{hex_width}x}")
                # Write in little-endian order (reverse)
                f.write(''.join(reversed(hex_values)) + '\n')
    
    print(f"Saved output to {filename}")


if __name__ == "__main__":
    # Configuration
    P_CH = 4
    N_CH = 16
    B_BIT = 16
    M_BIT = 32
    A_BIT = 8
    VEC_LEN = 64

    # Number of samples
    N_SAMPLES = VEC_LEN
    FOLD = N_CH // P_CH
    
    print("="*60)
    print("Configuration:")
    print(f"  P_CH={P_CH}, N_CH={N_CH}, FOLD={FOLD}")
    print(f"  B_BIT={B_BIT}, M_BIT={M_BIT}, A_BIT={A_BIT}")
    print(f"  VEC_LEN={VEC_LEN}, N_SAMPLES={N_SAMPLES}")
    print("="*60)
    
    # Generate test data
    # Input data: random values in range appropriate for B_BIT signed
    np.random.seed(42)
    input_data = np.random.randint(-128, 128, size=(N_SAMPLES, N_CH), dtype=np.int16)
    
    # Bias: random values in range appropriate for B_BIT signed
    bias = np.random.randint(-512, 512, size=N_CH, dtype=np.int16)
    
    # Mult: random values in range appropriate for M_BIT unsigned
    mult = np.random.randint(65536, 65536*4, size=N_CH, dtype=np.uint32)
    
    print("\nInput data shape:", input_data.shape)
    print("Bias shape:", bias.shape)
    print("Mult shape:", mult.shape)
    
    # Compute activation
    output = activation(input_data, bias, mult, B_BIT, A_BIT)
    
    print("\nOutput shape:", output.shape)
    print(f"Output range: [{output.min()}, {output.max()}]")
    
    # Save configuration to SystemVerilog header file
    print("\n" + "="*60)
    print("Saving data to files...")
    
    with open('actv_config.svh', 'w') as f:
        f.write("// Auto-generated configuration file for activation testbench\n")
        f.write("// Generated by gen_actv.py\n\n")
        f.write(f"localparam int unsigned P_CH = {P_CH};\n")
        f.write(f"localparam int unsigned N_CH = {N_CH};\n")
        f.write(f"localparam int unsigned B_BIT = {B_BIT};\n")
        f.write(f"localparam int unsigned M_BIT = {M_BIT};\n")
        f.write(f"localparam int unsigned A_BIT = {A_BIT};\n")
        f.write(f"localparam int unsigned VEC_LEN = {VEC_LEN};\n")
        f.write(f"\n// Derived parameters\n")
        f.write(f"localparam int unsigned FOLD = N_CH / P_CH;\n")
        f.write(f"localparam int unsigned N_SAMPLES = VEC_LEN * FOLD;\n")
    print("Saved configuration to actv_config.svh")
    
    # Save data files
    save_input_to_mem(input_data, 'actv_input.mem', 'w', P_CH, B_BIT)
    save_params_to_mem(bias, 'actv_bias.mem', 'w', N_CH, P_CH, B_BIT)
    save_params_to_mem(mult, 'actv_mult.mem', 'w', N_CH, P_CH, M_BIT)
    save_output_to_mem(output, 'actv_output.mem', 'w', P_CH, A_BIT)
    
    print("\n" + "="*60)
    print("Done!")
    print("="*60)